#include "nebomain.h"
#include "ui_nebomain.h"
#include <QFileDialog>
#include <QDebug>
#include <QFile>
#include <QMessageBox>
#include <QScrollArea>
#include <QScrollBar>
#include "utils.h"

inline void ConvertPath2PathStr(QString &path){
    path.replace("\\", "\\\\").replace("/", "\\\\");
}

NeboMain::NeboMain(QWidget *parent) :
    QMainWindow(parent),
    ui(new Ui::NeboMain)
{
    ui->setupUi(this);
    connect(ui->m_browse_btn, SIGNAL(clicked()), this, SLOT(browseRoamCapFile()));
    connect(ui->m_browse_btn_2, SIGNAL(clicked()), this, SLOT(browseDataRateCapFile()));

    connect(ui->m_calculate_btn, SIGNAL(clicked()), this, SLOT(calculateRoamTime()));
    connect(ui->m_calculate_btn_2, SIGNAL(clicked()), this, SLOT(calculateDataRate()));

    proc = new QProcess(this);
    init();
    setMACValidator();
    setAboutInfo();
    setFixedSize(width(), height());
    ui->m_log_scroll->ensureWidgetVisible( ui->m_log );
    QImage image(":/logo.jpg");
    if(!image.isNull())
        ui->mLogoImg->setPixmap(QPixmap::fromImage(image));


}
// Take nebo-input.cap that already renamed
// and call editcap and pass the filename and patch as parameters to the CLI command
bool NeboMain::TranslateCapFile()
{
    qDebug() << __FUNCTION__;
    // Check whether nebo-input.cap existed in directory or not
    QString cap_file = dirPath + QDir::separator()+ input_file_name;
    translatedCapFile = dirPath + QDir::separator() + translated_file_name;
    QString translate_file_str = translatedCapFile;
    QFile file(cap_file);
    if(!file.exists()){ // cap file does not exist
        QMessageBox::critical(this, tr("File error"),
                              tr("File nebo-input.cap does not exist in this directory\n"),
                              QMessageBox::Ok);
        return false;
    }
    // Create editcap command line
    ConvertPath2PathStr(cap_file);
    ConvertPath2PathStr(translate_file_str);
    QString command = "editcap -v -T  ieee-802-11-radiotap \"";
    command += cap_file + "\" \"";
    command += translate_file_str + "\"";

    qDebug() << command;
    proc->start(command);
    if(!proc->waitForFinished()){   // process error
        QMessageBox::critical(this, tr("Process error"),
                              tr("Could not execute editcap command\n"
                                 "Please check your environment."),
                              QMessageBox::Ok);
        return false;
    }
    qDebug("Success");
    return true;

}

// Take the input file wlan-da.cap/wlan-sa.cap and call wireshark
// and convert it to ASCII format based on measurement type
bool NeboMain::ConvertWlanCap2AsciiFile(MeasurementType type)
{
    qDebug() << __FUNCTION__;
    // Check whether nebo-input.cap existed in directory or not
    QString file_string = wlanFile;

    // Create ASCII format file
    asciiFile = dirPath + QDir::separator();
    asciiFile += (type == ROAM_TIME)?rt_ascii_file_name:dr_ascii_file_name;
    qDebug() << asciiFile;
    QString ascii_file_string = asciiFile;
    ConvertPath2PathStr(ascii_file_string);

    QFile file(file_string);
    if(!file.exists()){ // cap file does not exist
        QMessageBox::critical(this, tr("File error"),
                              tr("Wlan file does not exist\n"),
                              QMessageBox::Ok);
        return false;
    }
    // Create command for converting wlan file to ASCII format
    ConvertPath2PathStr(file_string);
    QString command = "cmd /c tshark.exe -r \"";
    command += file_string + "\" ";
    command += "-E separator=, -t ad -o wlan.enable_decryption:TRUE> \"";

    command += ascii_file_string + "\"";
    qDebug() << command;
    proc->start(command);
    if(!proc->waitForFinished()){ // process error
        QMessageBox::critical(this, tr("Process error"),
                              tr("Could not execute tshark command\n"
                                 "Please check your environment."),
                              QMessageBox::Ok);
        return false;
    }

    qDebug("Success");
    return true;
}

NeboMain::~NeboMain()
{
    delete ui;
}
// Get MAC address from UI. MAC address will be used for
// filtering translated capture file
inline QString NeboMain::getMACAddress(const MeasurementType type)
{
    qDebug() << __FUNCTION__;
    QString macAddress;
    if(type == ROAM_TIME){
        macAddress = ui->m_ROAM_MAC_add1->text()+ ":";
        macAddress += ui->m_ROAM_MAC_add2->text()+ ":";
        macAddress += ui->m_ROAM_MAC_add3->text()+ ":";
        macAddress += ui->m_ROAM_MAC_add4->text()+ ":";
        macAddress += ui->m_ROAM_MAC_add5->text()+ ":";
        macAddress += ui->m_ROAM_MAC_add6->text();
    }else{
        macAddress = ui->m_DR_MAC_add1->text()+ ":";
        macAddress += ui->m_DR_MAC_add2->text()+ ":";
        macAddress += ui->m_DR_MAC_add3->text()+ ":";
        macAddress += ui->m_DR_MAC_add4->text()+ ":";
        macAddress += ui->m_DR_MAC_add5->text()+ ":";
        macAddress += ui->m_DR_MAC_add6->text();
    }

    qDebug() << "MAC address:" << macAddress;
    return macAddress;
}

// Call external command to filter translated capture file to wlan capture file.
// DOS command line: tshark.exe –r (path)translated.cap -w (path)wlan-da.cap  -R "wlan.da == MAC”
bool NeboMain::filterTranslatedCap(const QString mac_add, const MeasurementType type)
{
    qDebug() << __FUNCTION__;
    // Create wlan file path
    wlanFile = dirPath + QDir::separator();
    wlanFile += (type == ROAM_TIME)?rt_wlan_file_name:dr_wlan_file_name;
    qDebug() << wlanFile;
    QString wlan_file_str = wlanFile;
    ConvertPath2PathStr(wlan_file_str);
    QString translate_file_str = translatedCapFile;
    ConvertPath2PathStr(translate_file_str);

    // Create filter command
    QString command ="tshark.exe -r \"" + translate_file_str + "\"";
    command += " -w \""+ wlan_file_str + "\"";
    command += (type == ROAM_TIME)?" -R \"wlan.da==":" -R \"wlan.sa==";
    command += mac_add+ "\"";
    qDebug() << command;
    proc->start(command);
    if(!proc->waitForFinished()){
        qDebug() << "Filter process has failed !";
        QMessageBox::critical(this, tr("Process error"),
                              tr("Could not execute filtering wlan.\n"
                                 "Please check your environment"), QMessageBox::Ok);
        return false;
    }
    return true;

}

// Parse string line in ASCII file to find MAC address
// by using regular expression
QString NeboMain::parseMACAddress(const QString str)
{
    qDebug() << __FUNCTION__;
    // Set MAC address pattern in regular expression
    QRegExp rx("([0-9a-fA-F][0-9a-fA-F]:){5}([0-9a-fA-F][0-9a-fA-F])");
    QString mac_add = "";
    int pos = 0;

    while ((pos = rx.indexIn(str, pos)) != -1) {
        mac_add = rx.cap();
        break;
    }
    return mac_add;

}

// Read and analyze ASCII file to calculate
// roam time or data rate depending on measurement type
void NeboMain::readAsciiFile(MeasurementType type)
{
    qDebug() << __FUNCTION__;
    qDebug() << asciiFile;
    QFile file(asciiFile);
    // Check whether ASCII file existed
    if(!file.exists()){ // file does not exist
        QMessageBox::critical(this, tr("File error"),
                              tr("ASCII file does not exist."),
                              QMessageBox::Ok);
        return;
    }
    QString str;
    QString mac_add;
    QString replace_str = "";
    if (file.open(QIODevice::ReadOnly | QIODevice::Text))   // Open file for read only
    {
        QTextStream in(&file);
        while (!in.atEnd()) {
            str = in.readLine();
            mac_add = parseMACAddress(str);
            if(mac_add.isEmpty())
                continue;
            if(macAddress != mac_add){
                replace_str = "<font color=\"Red\">" + mac_add + "</font>";
                str.replace(mac_add, replace_str);
            }
            str += "<br>";
            if(type == ROAM_TIME){
                if(str.contains("Reassociation Response")){ // calculate roam after reassociation response
                    if(in.atEnd())
                        continue;
                    QString temp = in.readLine();
                    QString mac1 = parseMACAddress(temp);
                    if(mac1 == macAddress && mac1 != "")
                        continue;
                    QString absolute_time1 = parseAbsoluteTime(temp);
                    temp = in.readLine();
                    QString mac2 = parseMACAddress(temp);
                    if(mac2 == macAddress && mac2 != "")
                        continue;
                    QString absolute_time2 = parseAbsoluteTime(temp);
                    QStringList list1 = absolute_time1.split(":");
                    QStringList list2 = absolute_time2.split(":");
                    int size1 = list1.size();
                    int size2 = list2.size();
                    if((size1 != size2) && size1 != 3 && size2 != 3){
                        qDebug() << "Error in determining absolute time";
                        continue;
                    }

                    float roam_time = abs(list1.at(0).toInt() - list2.at(0).toInt())*10000;
                    roam_time += abs(list1.at(1).toInt() - list2.at(1).toInt())*100;
                    roam_time += abs(list1.at(2).toFloat() - list2.at(2).toFloat());
                    roamCount++;
                    roamSum += roam_time;
                    if(roam_time > longestRoamTime)
                        longestRoamTime = roam_time;
                    qDebug() << "Roam time:" << roam_time << " ms";

                }
                //ui->m_log->insertHtml(str);
            }else{
                parseDataRate(str);
                //ui->m_log_2->insertHtml(str);
            }

        }
        if(type == ROAM_TIME){
            QString result = "";
            if(roamCount){
                qDebug() << "Average roam time - " << roamSum/roamCount;
                qDebug() << "Longest roam time - " << longestRoamTime;
                result = "<font color=\"Red\"> Results..<br>Average roam time - ";
                result.append(QString("%1").arg(roamSum/roamCount));
                result += " ms<br>Longest roam time - ";
                result.append(QString("%1").arg(longestRoamTime));
                result += " ms<br>Completed.</font>";
            }else{
                result = "<font color=\"Red\">No roams detected.</font>";
            }

            ui->m_log->insertHtml(result);
        }
        else{
            //qDebug() << "Average data rate - " << averageDr/drCount;
             QString result;
            if(drCount){
                result = "<font color=\"Red\">Results..<br>Average data rate - ";
                result.append(QString("%1").arg(averageDr/drCount));
                result += " Mbps<br>Completed.</font>";
            }else{
                result = "<font color=\"Red\">Could not calculate data rate </font>";
            }


            ui->m_log_2->insertHtml(result);

        }

    }
}

// Parse absolute time that is used for calculating roam time
QString NeboMain::parseAbsoluteTime(const QString str)
{

    QRegExp regexp("([0-9][0-9]:){2}([0-9]*\.?[0-9]+)");
    QString absolute_time = "";
    int pos = 0;

    while ((pos = regexp.indexIn(str, pos)) != -1) {
        absolute_time = regexp.cap();
        break;
    }
    qDebug() << absolute_time;
    return absolute_time;
}

void NeboMain::parseDataRate(const QString str)
{
    QRegExp regexp("\\d+\\.\\d+\\s802.11");
    QString temp = "";
    int pos = 0;

    while ((pos = regexp.indexIn(str, pos)) != -1) {
        temp = regexp.cap();
        break;
    }
    qDebug() << "Data rate string:" << temp;
    if(!temp.isEmpty()){
        QStringList list = temp.split(" ");
        float data_rate = list.at(0).toFloat();
        qDebug() << "Data rate:" << data_rate;
        if(data_rate > 0.0f){
            averageDr += data_rate;
            drCount++;
        }
    }

}

// Set MAC address validation for MAC textboxes.
// It only accepts  number characters (0-9) or
// characters from a-f or A-F
void NeboMain::setMACValidator()
{
    QRegExp regexp("[0-9a-fA-F][0-9a-fA-F]");
    QRegExpValidator *mac_validator = new QRegExpValidator(regexp, this);
    // MAC validator for roam time
    ui->m_ROAM_MAC_add1->setValidator(mac_validator);
    ui->m_ROAM_MAC_add2->setValidator(mac_validator);
    ui->m_ROAM_MAC_add3->setValidator(mac_validator);
    ui->m_ROAM_MAC_add4->setValidator(mac_validator);
    ui->m_ROAM_MAC_add5->setValidator(mac_validator);
    ui->m_ROAM_MAC_add6->setValidator(mac_validator);
    // MAC validator for data rate
    ui->m_DR_MAC_add1->setValidator(mac_validator);
    ui->m_DR_MAC_add2->setValidator(mac_validator);
    ui->m_DR_MAC_add3->setValidator(mac_validator);
    ui->m_DR_MAC_add4->setValidator(mac_validator);
    ui->m_DR_MAC_add5->setValidator(mac_validator);
    ui->m_DR_MAC_add6->setValidator(mac_validator);

}

bool NeboMain::checkEmptyMACAddr(MeasurementType meas_type)
{
    bool is_empty = false;
    if(meas_type == ROAM_TIME){
        if(ui->m_ROAM_MAC_add1->text().isEmpty() ||
           ui->m_ROAM_MAC_add2->text().isEmpty() ||
           ui->m_ROAM_MAC_add3->text().isEmpty() ||
           ui->m_ROAM_MAC_add4->text().isEmpty() ||
           ui->m_ROAM_MAC_add5->text().isEmpty() ||
           ui->m_ROAM_MAC_add6->text().isEmpty()){
            is_empty = true;
        }


    }else{
        if(ui->m_DR_MAC_add1->text().isEmpty() ||
           ui->m_DR_MAC_add2->text().isEmpty() ||
           ui->m_DR_MAC_add3->text().isEmpty() ||
           ui->m_DR_MAC_add4->text().isEmpty() ||
           ui->m_DR_MAC_add5->text().isEmpty() ||
           ui->m_DR_MAC_add6->text().isEmpty()){
            is_empty = true;
        }

    }

    if(is_empty){
        QMessageBox::critical(this, tr("MAC address error"),
                              tr("MAC address is not valid\n"
                                 "Please check it again!"),
                              QMessageBox::Ok);
        return false;
    }
    return true;
}

//Create uniformity of extension types by renaming non .cap files to .cap…
bool NeboMain::renameNonCapFile()
{
    QString inputCapFileStr = inputCapFile;
    ConvertPath2PathStr(inputCapFileStr);
    // Create command for renaming non-cap file to input-cap.cap
    QString command = "cmd /c rename \"";
    command += inputCapFileStr +"\" ";
    command += input_file_name;
    qDebug() << command;
    proc->start(command);
    if(!proc->waitForFinished()){ // the process will be fail if executing command meets error
        // Pop up error message dialog for user
        QMessageBox::critical(this, tr("Process error"),
                              tr("Could not execute Rename DOS command."),
                              QMessageBox::Ok);
        return false;
    }
    return true;
}

void NeboMain::browseRoamCapFile()
{
    // Set filename to textbox
    browseFile();
    ui->m_file_path->setText(inputCapFile);
}

void NeboMain::browseDataRateCapFile()
{
    // Set filename to textbox
    browseFile();
    ui->m_file_path_2->setText(inputCapFile);
}

void NeboMain::calculateRoamTime()
{
    qDebug() << "In " << __FUNCTION__;
    // Check empty MAC address
    ui->m_log->clear();

    if(checkEmptyMACAddr(ROAM_TIME) == false)
        return;
    macAddress = getMACAddress(ROAM_TIME);
    QString string = "<font color=\"Red\"> Calculating roam time for MAC:";
    string += macAddress + "...<br></font>";
    ui->m_log->insertHtml(string);

    if(renameNonCapFile() == false)
        return;
    if(TranslateCapFile() == false)
        return;

    if(filterTranslatedCap(macAddress, ROAM_TIME) == false)
        return;
    if(ConvertWlanCap2AsciiFile(ROAM_TIME) == false)
        return;

//    ascii_file = "E:/ascii-cap.txt";
    readAsciiFile(ROAM_TIME);
    qDebug() << "End";
}

void NeboMain::calculateDataRate()
{
    qDebug() << "In " << __FUNCTION__;
    // Check empty MAC address
    if(checkEmptyMACAddr(DATA_RATE) == false)
        return;
    macAddress = getMACAddress(DATA_RATE);
    QString string = "<font color=\"Red\"> Calculating average data rate for MAC:";
    string += macAddress + "...<br></font>";
    ui->m_log_2->insertHtml(string);

    if(renameNonCapFile() == false)
        return;
    if(TranslateCapFile() == false)
        return;
    if(filterTranslatedCap(macAddress, DATA_RATE) == false)
        return;
    if(ConvertWlanCap2AsciiFile(DATA_RATE) == false)
        return;

    readAsciiFile(DATA_RATE);
    qDebug() << "End";
}

void NeboMain::browseFile()
{
    inputCapFile = QFileDialog::getOpenFileName(this,
        tr("Open capture file"), QDir::homePath(),
        tr("Capture Files (*.cap *.pcap *.apc)"));  // only accept .cap, .pcap, .apc
    QDir dir = QFileInfo(inputCapFile).absoluteDir(); // get directory that contains the file
    dirPath = dir.absolutePath();
}

void NeboMain::setAboutInfo()
{
    QString about_str = "Nebo is completely free to use and distribute. ";
    about_str.append("All wireless captures of mobile clients ");
    about_str.append("*MUST* contain the MAC address of a client that is already associated to the network. ");
    about_str.append("Any deviation from this rule will yield results that are skewed. ");
    about_str.append("Please submit any comments or suggestions for this tool to ");
    about_str.append("<a href=\"mailto:Joypher@wifinut.net\">Joypher@wifinut.net </a>");
    ui->m_about->insertHtml(about_str);
}

